from __future__ import print_function, division

import torch

import seq2seq
from seq2seq.loss import NLLLoss, loss_fuc
from torch.nn import CrossEntropyLoss

class Evaluator(object):
    """ Class to evaluate models with given datasets.

    Args:
        loss (seq2seq.loss, optional): loss for evaluator (default: seq2seq.loss.NLLLoss)
        batch_size (int, optional): batch size for evaluator (default: 64)
    """

    def __init__(self, loss=NLLLoss(), model_name='MoeTransformer', batch_size=64, device=None):
        self.loss = loss
        self.batch_size = batch_size
        self.device = device
        self.vocab_size = 40000
        self.model_name = model_name

    def evaluate(self, model, data):
        """ Evaluate a model on given dataset and return performance.

        Args:
            model (seq2seq.models): model to evaluate
            data (seq2seq.dataset.dataset.Dataset): dataset to evaluate against

        Returns:
            loss (float): loss of the given model on the given dataset
        """
        model.eval()

        loss = self.loss
        device = self.device

        loss.reset()
        match = 0
        total = 0

        tgt_vocab = data.dataset.tgt_vocab
        pad = tgt_vocab.word2idx[tgt_vocab.pad_token]

        with torch.no_grad():
            for batch in data:
                input_variable = batch['src'].to(device)
                target_variable = batch['tgt'].to(device)
                src_lens = batch['src_len'].view(-1).to(device)
                tgt_lens = batch['tgt_len'].view(-1).to(device)

                if self.model_name == 'MoeTransformer':
                    his_variable = batch['his'].to(device)
                    his_seg_variable = batch['his_seg'].to(device)
                    decoder_outputs = model(input_variable, his_variable, his_seg_variable, target_variable[:, :-1])
                else:
                    decoder_outputs = model(input_variable , target_variable[:, :-1])
                loss, ppl = loss_fuc(decoder_outputs, target_variable[:, 1:])

                b, t, d = decoder_outputs.size()
                for step, step_output in enumerate(decoder_outputs.transpose(0,1)):
                    step_target =  target_variable[:,step+1].contiguous().view(-1)
                    mask = step_target.ne(1)
                    correct = step_output.argmax(1).view(-1).eq(step_target).masked_select(mask).sum().item()
                    match += correct
                    total += mask.sum().item()
                # Evaluation
                # seqlist = other['sequence']
                # for step, step_output in enumerate(decoder_outputs):
                #     batch_size = target_variable.size(1)
                #     flattened_output = step_output.view(-1, self.vocab_size) 
                #     # step_target = target_variable[:,step+1].contiguous().view(-1)
                #     step_target = target_variable[step+1,:].contiguous().view(-1)
                #     mask = step_target.ne(1)
                #     masked_target = step_target.masked_select(mask)
                #     output_mask = mask.unsqueeze(1).expand(mask.size(0), self.vocab_size)
                #     masked_output = flattened_output.masked_select(output_mask).view(-1, self.vocab_size)
                #     loss.eval_batch(masked_output, masked_target)

                #     correct = decoder_outputs[step].argmax(1).view(-1).eq(step_target).masked_select(mask).sum().item()
                #     match += correct
                #     total += mask.sum().item()
        if total == 0:
            accuracy = float('nan')
        else:
            accuracy = match / total

        return ppl, accuracy
